{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How-to Finetune\n",
    "\n",
    "This tutorial shows how to adapt a pretrained model to a different, eventually much smaller dataset, a concept called finetuning. Finetuning is well-established in machine learning and thus nothing new. Generally speaking, the idea is to use a (very) large and diverse dataset to learn a general understanding of the underlying problem first and then, in a second step, adapt this general model to the target data. Usually, especially if the available target data is limited, pretraining plus finetuning yields (much) better results than only considering the final target data. \n",
    "\n",
    "The connection to hydrology is the following: Often, researchers or operators are only interested in a single basin. However, considering that a Deep Learning (DL) model has to learn all (physical) process understanding from the available training data, it might be understandable that the data records of a single basin might not be enough (see e.g. the presentation linked at [this](https://meetingorganizer.copernicus.org/EGU2020/EGU2020-8855.html) EGU'20 abstract)\n",
    "\n",
    "This is were we apply the concept of pretraining and finetuning: First, we train a DL model (e.g. an LSTM) with a large and diverse, multi-basin dataset (e.g. CAMELS) and then finetune this model to our basin of interest. Everything you need is available in the `neuralHydrology` package and in this notebook we will give you an overview of how to actually do it.\n",
    "\n",
    "**Note**: Finetuning can be a tedious task and is usually very sensitive to the learning rate as well as the number of epochs used for finetuning. One reason is that the pretrained models are usually quite large. In fact, most often they are much larger than what would be possible to train for just a single basin. So during finetuning, we have to make sure that this large capacity is not negatively impacting our model results. Common approaches are to a) only allow parts of the model to be adapted during finetuning and/or b) to train with a much lower learning rate. So far, no publication was published that presents a universally working approach for finetuning in hydrology. So be aware that the results may vary and you might need to invest some time before finding a good strategy. However, in our experience it was always possible to get better results _with_ finetuning than without.\n",
    "\n",
    "**To summarize**: If you are interested in getting the best-performing Deep Learning model for a single basin, pretraining on a large and diverse dataset, followed by finetuning the pretrained model on your target basin is the way to go."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import pickle\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from neuralhydrology.nh_run import start_run, eval_run, finetune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pretraining\n",
    "\n",
    "In the first step, we need to pretrain our model on a large and possibly diverse dataset. Our target basin does not necessarily have to be a part of this dataset, but usually it should be better to include it.\n",
    "\n",
    "For the sake of the demonstration, we will train an LSTM on the CAMELS US dataset and then finetune this model to a random basin. Note that it is possible to use other inputs during pretraining and finetuning, if additional embedding layers (before the LSTM) are used, which we will ignore for now. Furthermore, we will concentrate only on demonstrating the \"how-to\" rather than striving for best-possible performance. To save time and energy, we will only pretrain the model for a small number of epochs. When striving for the best possible performance, you should make sure that you pretrain the model as best as possible, before starting to finetune.\n",
    "\n",
    "We will stick closely to the model and experimental setup from [Kratzert et al. (2019)](https://hess.copernicus.org/articles/23/5089/2019/hess-23-5089-2019.html). To summarize:\n",
    "- A single LSTM layer with a hidden size of 128.\n",
    "- Input sequences are 365 days and the prediction is made at the last timestep.\n",
    "- For the sake of this demonstration, we will only consider the 5 meteorological variables from the [extended Maurer](https://doi.org/10.4211/hs.17c896843cf940339c3c3496d0c1c077) forcing data. Either download these forcings and place the `maurer_extended` folder into the `basin_mean_forcing` folder within the CAMELS US root directory or change the forcing product and dynamic inputs in the config file.\n",
    "- We will use the same CAMELS attributes, as in the publication mentioned above, as additional inputs at every time step so that the model can learn different hydrological behaviors depending on the catchment properties.\n",
    "\n",
    "For more details, take a look at the config print-out below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-12 21:30:10,861: Logging to /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/output.log initialized.\n",
      "2021-01-12 21:30:10,862: ### Folder structure created at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010\n",
      "2021-01-12 21:30:10,863: ### Run configurations for cudalstm_maurer_531_basins\n",
      "2021-01-12 21:30:10,863: experiment_name: cudalstm_maurer_531_basins\n",
      "2021-01-12 21:30:10,864: run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010\n",
      "2021-01-12 21:30:10,865: train_basin_file: 531_basin_list.txt\n",
      "2021-01-12 21:30:10,866: validation_basin_file: 531_basin_list.txt\n",
      "2021-01-12 21:30:10,867: test_basin_file: 531_basin_list.txt\n",
      "2021-01-12 21:30:10,868: train_start_date: 1999-10-01 00:00:00\n",
      "2021-01-12 21:30:10,868: train_end_date: 2008-09-30 00:00:00\n",
      "2021-01-12 21:30:10,869: validation_start_date: 1980-10-01 00:00:00\n",
      "2021-01-12 21:30:10,870: validation_end_date: 1989-09-30 00:00:00\n",
      "2021-01-12 21:30:10,871: test_start_date: 1989-10-01 00:00:00\n",
      "2021-01-12 21:30:10,875: test_end_date: 1999-09-30 00:00:00\n",
      "2021-01-12 21:30:10,876: seed: 123\n",
      "2021-01-12 21:30:10,877: device: cuda:0\n",
      "2021-01-12 21:30:10,878: validate_every: 1\n",
      "2021-01-12 21:30:10,878: validate_n_random_basins: 531\n",
      "2021-01-12 21:30:10,879: metrics: ['NSE']\n",
      "2021-01-12 21:30:10,880: model: cudalstm\n",
      "2021-01-12 21:30:10,881: head: regression\n",
      "2021-01-12 21:30:10,882: hidden_size: 128\n",
      "2021-01-12 21:30:10,883: initial_forget_bias: 3\n",
      "2021-01-12 21:30:10,887: output_dropout: 0.4\n",
      "2021-01-12 21:30:10,887: output_activation: linear\n",
      "2021-01-12 21:30:10,888: optimizer: Adam\n",
      "2021-01-12 21:30:10,889: loss: NSE\n",
      "2021-01-12 21:30:10,890: learning_rate: {0: 0.001, 1: 0.0005}\n",
      "2021-01-12 21:30:10,891: batch_size: 256\n",
      "2021-01-12 21:30:10,892: epochs: 3\n",
      "2021-01-12 21:30:10,893: clip_gradient_norm: 1\n",
      "2021-01-12 21:30:10,893: predict_last_n: 1\n",
      "2021-01-12 21:30:10,894: seq_length: 365\n",
      "2021-01-12 21:30:10,895: num_workers: 8\n",
      "2021-01-12 21:30:10,896: log_interval: 5\n",
      "2021-01-12 21:30:10,896: log_tensorboard: True\n",
      "2021-01-12 21:30:10,897: save_weights_every: 1\n",
      "2021-01-12 21:30:10,898: save_validation_results: True\n",
      "2021-01-12 21:30:10,899: dataset: camels_us\n",
      "2021-01-12 21:30:10,900: data_dir: /data/Hydrology/CAMELS_US\n",
      "2021-01-12 21:30:10,900: forcings: ['maurer_extended']\n",
      "2021-01-12 21:30:10,901: dynamic_inputs: ['prcp(mm/day)', 'srad(W/m2)', 'tmax(C)', 'tmin(C)', 'vp(Pa)']\n",
      "2021-01-12 21:30:10,902: target_variables: ['QObs(mm/d)']\n",
      "2021-01-12 21:30:10,903: static_attributes: ['elev_mean', 'slope_mean', 'area_gages2', 'frac_forest', 'lai_max', 'lai_diff', 'gvf_max', 'gvf_diff', 'soil_depth_pelletier', 'soil_depth_statsgo', 'soil_porosity', 'soil_conductivity', 'max_water_content', 'sand_frac', 'silt_frac', 'clay_frac', 'carbonate_rocks_frac', 'geol_permeability', 'p_mean', 'pet_mean', 'aridity', 'frac_snow', 'high_prec_freq', 'high_prec_dur', 'low_prec_freq', 'low_prec_dur']\n",
      "2021-01-12 21:30:10,903: number_of_basins: 531\n",
      "2021-01-12 21:30:10,904: train_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/train_data\n",
      "2021-01-12 21:30:10,904: img_log_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/img_log\n",
      "2021-01-12 21:30:10,943: ### Device cuda:0 will be used for training\n",
      "2021-01-12 21:30:13,667: Loading basin data into xarray data set.\n",
      "100%|██████████| 531/531 [00:35<00:00, 14.80it/s]\n",
      "2021-01-12 21:30:49,952: Calculating target variable stds per basin\n",
      "100%|██████████| 531/531 [00:00<00:00, 827.94it/s]\n",
      "2021-01-12 21:30:50,709: Create lookup table and convert to pytorch tensor\n",
      "100%|██████████| 531/531 [00:15<00:00, 35.16it/s]\n",
      "2021-01-12 21:31:06,316: Setting learning rate to 0.0005\n",
      "# Epoch 1: 100%|██████████| 6821/6821 [03:49<00:00, 29.66it/s, Loss: 0.0039]\n",
      "2021-01-12 21:34:56,319: Epoch 1 average loss: 0.03816824801886893\n",
      "# Validation: 100%|██████████| 531/531 [03:30<00:00,  2.52it/s]\n",
      "2021-01-12 21:38:26,862: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/validation/model_epoch001/validation_results.p\n",
      "2021-01-12 21:38:26,870:  -- Median validation metrics:NSE: 0.63335\n",
      "# Epoch 2: 100%|██████████| 6821/6821 [03:54<00:00, 29.08it/s, Loss: 0.0035]\n",
      "2021-01-12 21:42:21,468: Epoch 2 average loss: 0.024925640246485033\n",
      "# Validation: 100%|██████████| 531/531 [02:32<00:00,  3.48it/s]\n",
      "2021-01-12 21:44:54,314: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/validation/model_epoch002/validation_results.p\n",
      "2021-01-12 21:44:54,321:  -- Median validation metrics:NSE: 0.65184\n",
      "# Epoch 3: 100%|██████████| 6821/6821 [03:55<00:00, 28.91it/s, Loss: 0.0006]\n",
      "2021-01-12 21:48:50,226: Epoch 3 average loss: 0.022466508427862103\n",
      "# Validation: 100%|██████████| 531/531 [02:29<00:00,  3.55it/s]\n",
      "2021-01-12 21:51:20,057: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/validation/model_epoch003/validation_results.p\n",
      "2021-01-12 21:51:20,065:  -- Median validation metrics:NSE: 0.68917\n"
     ]
    }
   ],
   "source": [
    "config_file = Path(\"531_basins.yml\")\n",
    "start_run(config_file=config_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We end with an okay'ish model that should be enough for the purpose of this demonstration. Remember we only train for a limited number of epochs here.\n",
    "\n",
    "Next, let's look in the `runs/` folder, where the folder of this model is stored to lookup the exact name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cudalstm_maurer_531_basins_1201_213010\r\n"
     ]
    }
   ],
   "source": [
    "!ls runs/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll load the validation results into memory so we can select a basin to demonstrate how to finetune based on the model performance. Here, we will select a random basin from the lower 50% of the NSE distribution, i.e. a basin where the NSE is below the median NSE. Usually, you'll see better performance gains for basins with lower model performance than for those where the base model is already really good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Median NSE of the validation period 0.689\n",
      "Selected basin: 04122200 with an NSE of 0.629\n"
     ]
    }
   ],
   "source": [
    "# Load validation results for the last epoch\n",
    "run_dir = Path(\"runs/cudalstm_maurer_531_basins_1201_213010\")\n",
    "with open(run_dir / \"validation\" / \"model_epoch003\" / \"validation_results.p\", \"rb\") as fp:\n",
    "    validation_results = pickle.load(fp)\n",
    "\n",
    "# Compute the median NSE from all basins, where discharge observations are available for that period\n",
    "median_nse = np.median([v[\"1D\"][\"NSE\"] for v in validation_results.values() if \"NSE\" in v[\"1D\"].keys()])\n",
    "print(f\"Median NSE of the validation period {median_nse:.3f}\")\n",
    "\n",
    "# Select a random basins from the lower 50% of the NSE distribution\n",
    "basins = []\n",
    "for k, v in validation_results.items():\n",
    "    if (\"NSE\" in v[\"1D\"].keys()) and (v[\"1D\"][\"NSE\"] < median_nse):\n",
    "        basins.append(k)\n",
    "basin = np.random.choice(basins)\n",
    "\n",
    "print(f\"Selected basin: {basin} with an NSE of {validation_results[basin]['1D']['NSE']:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finetuning\n",
    "\n",
    "Next, we will show how to perform finetuning for the basin selected above, based on the model we just trained. The function to use is `finetune` from `neuralhydrology.nh_run` if you want to train from within a script or notebook. If you want to start finetuning from the command line, you can also call the `nh-run` utility with the `finetune` argument, instead of e.g. `train` or `evaluate`.\n",
    "\n",
    "The only thing required, similar to the model training itself, is a config file. This config however has slightly different requirements to a normal model config and works slightly different:\n",
    "- The config has to contain the following two arguments:\n",
    "    - `base_run_dir`: The path to the directory of the pre-trained model.\n",
    "    - `finetune_modules`: Which parts of the pre-trained model you want to finetune. Check the documentation of each model class for a list of all possible parts. Often only parts, e.g. the output layer, are trained during finetuning and the rest is kept fixed. There is no general rule of thumb and most likely you will have to try both.\n",
    "- Any additional argument contained in this config will overwrite the config argument of the pre-trained model. Everything _not_ specified will be taken from the pre-trained model. That is, you can e.g. specify a new basin file in the finetuning config (by `train_basin_file`) to finetune the pre-trained model on a different set of basins, or even just a single basin as we will do in this notebook. You can also change the learning rate, loss function, evaluation metrics and so on. The only thing you can not change are arguments that change the model architecture (e.g. `model`, `hidden_size` etc.), because this leads to errors when you try to load the pre-trained weights into the initialized model.\n",
    "\n",
    "Let's have a look at the `finetune.yml` config that we prepared for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# --- Experiment configurations --------------------------------------------------------------------\r\n",
      "\r\n",
      "# experiment name, used as folder name\r\n",
      "experiment_name: cudalstm_maurer_531_basins_finetuned\r\n",
      "\r\n",
      "# files to specify training, validation and test basins (relative to code root or absolute path)\r\n",
      "train_basin_file: finetune_basin.txt\r\n",
      "validation_basin_file: finetune_basin.txt\r\n",
      "test_basin_file: finetune_basin.txt\r\n",
      "\r\n",
      "# --- Training configuration -----------------------------------------------------------------------\r\n",
      "\r\n",
      "# specify learning rates to use starting at specific epochs (0 is the initial learning rate)\r\n",
      "learning_rate:\r\n",
      "    0: 5e-4\r\n",
      "    2: 5e-5\t\r\n",
      "\r\n",
      "# Number of training epochs\r\n",
      "epochs: 10\r\n",
      "\r\n",
      "finetune_modules:\r\n",
      "- head\r\n",
      "- lstm\r\n"
     ]
    }
   ],
   "source": [
    "!cat finetune.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So out of the two arguments that are required, `base_run_dir` is still missing. We will add the argument from here and point at the directory of the model we just trained. Furthermore, we point to a new file for training, validation and testing, called `finetune_basin.txt`, which does not yet exist. We will create this file and add the basin we selected above as the only basin we want to use here. The rest are some changes to the learning rate and the number of training epochs as well as a new name. Also note that here, we train the full model, by selecting all model parts available for the `CudaLSTM` under `finetune_modules`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the path to the pre-trained model to the finetune config\n",
    "with open(\"finetune.yml\", \"a\") as fp:\n",
    "    fp.write(f\"base_run_dir: {run_dir.absolute()}\")\n",
    "    \n",
    "# Create a basin file with the basin we selected above\n",
    "with open(\"finetune_basin.txt\", \"w\") as fp:\n",
    "    fp.write(basin)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With that, we are ready to start the finetuning. As mentioned above, we have two options to start finetuning:\n",
    "1. Call the `finetune()` function from a different Python script or a Jupyter Notebook with the path to the config.\n",
    "2. Start the finetuning from the command line by calling\n",
    "\n",
    "```bash\n",
    "nh-run finetune --config-file /path/to/config.yml\n",
    "```\n",
    "\n",
    "Here, we will use the first option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-12 21:55:00,954: Logging to /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/output.log initialized.\n",
      "2021-01-12 21:55:00,956: ### Folder structure created at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500\n",
      "2021-01-12 21:55:00,957: ### Start finetuning with pretrained model stored in /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010\n",
      "2021-01-12 21:55:00,958: ### Run configurations for cudalstm_maurer_531_basins_finetuned\n",
      "2021-01-12 21:55:00,959: batch_size: 256\n",
      "2021-01-12 21:55:00,960: clip_gradient_norm: 1\n",
      "2021-01-12 21:55:00,963: commit_hash: c271d04\n",
      "2021-01-12 21:55:00,964: data_dir: /data/Hydrology/CAMELS_US\n",
      "2021-01-12 21:55:00,966: dataset: camels_us\n",
      "2021-01-12 21:55:00,967: device: cuda:0\n",
      "2021-01-12 21:55:00,967: dynamic_inputs: ['prcp(mm/day)', 'srad(W/m2)', 'tmax(C)', 'tmin(C)', 'vp(Pa)']\n",
      "2021-01-12 21:55:00,968: epochs: 10\n",
      "2021-01-12 21:55:00,970: experiment_name: cudalstm_maurer_531_basins_finetuned\n",
      "2021-01-12 21:55:00,971: forcings: ['maurer_extended']\n",
      "2021-01-12 21:55:00,971: head: regression\n",
      "2021-01-12 21:55:00,972: hidden_size: 128\n",
      "2021-01-12 21:55:00,972: img_log_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/img_log\n",
      "2021-01-12 21:55:00,974: initial_forget_bias: 3\n",
      "2021-01-12 21:55:00,975: learning_rate: {0: 0.0005, 2: 5e-05}\n",
      "2021-01-12 21:55:00,975: log_interval: 5\n",
      "2021-01-12 21:55:00,976: log_tensorboard: True\n",
      "2021-01-12 21:55:00,977: loss: NSE\n",
      "2021-01-12 21:55:00,978: metrics: ['NSE']\n",
      "2021-01-12 21:55:00,979: model: cudalstm\n",
      "2021-01-12 21:55:00,980: num_workers: 8\n",
      "2021-01-12 21:55:00,980: number_of_basins: 1\n",
      "2021-01-12 21:55:00,981: optimizer: Adam\n",
      "2021-01-12 21:55:00,982: output_activation: linear\n",
      "2021-01-12 21:55:00,982: output_dropout: 0.4\n",
      "2021-01-12 21:55:00,984: package_version: 0.9.8-beta\n",
      "2021-01-12 21:55:00,984: predict_last_n: 1\n",
      "2021-01-12 21:55:00,985: run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500\n",
      "2021-01-12 21:55:00,986: save_validation_results: True\n",
      "2021-01-12 21:55:00,987: save_weights_every: 1\n",
      "2021-01-12 21:55:00,987: seed: 123\n",
      "2021-01-12 21:55:00,988: seq_length: 365\n",
      "2021-01-12 21:55:00,989: static_attributes: ['elev_mean', 'slope_mean', 'area_gages2', 'frac_forest', 'lai_max', 'lai_diff', 'gvf_max', 'gvf_diff', 'soil_depth_pelletier', 'soil_depth_statsgo', 'soil_porosity', 'soil_conductivity', 'max_water_content', 'sand_frac', 'silt_frac', 'clay_frac', 'carbonate_rocks_frac', 'geol_permeability', 'p_mean', 'pet_mean', 'aridity', 'frac_snow', 'high_prec_freq', 'high_prec_dur', 'low_prec_freq', 'low_prec_dur']\n",
      "2021-01-12 21:55:00,990: target_variables: ['QObs(mm/d)']\n",
      "2021-01-12 21:55:00,990: test_basin_file: finetune_basin.txt\n",
      "2021-01-12 21:55:00,991: test_end_date: 1999-09-30 00:00:00\n",
      "2021-01-12 21:55:00,991: test_start_date: 1989-10-01 00:00:00\n",
      "2021-01-12 21:55:00,992: train_basin_file: finetune_basin.txt\n",
      "2021-01-12 21:55:00,992: train_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/train_data\n",
      "2021-01-12 21:55:00,993: train_end_date: 2008-09-30 00:00:00\n",
      "2021-01-12 21:55:00,993: train_start_date: 1999-10-01 00:00:00\n",
      "2021-01-12 21:55:00,993: validate_every: 1\n",
      "2021-01-12 21:55:00,995: validate_n_random_basins: 531\n",
      "2021-01-12 21:55:00,995: validation_basin_file: finetune_basin.txt\n",
      "2021-01-12 21:55:00,996: validation_end_date: 1989-09-30 00:00:00\n",
      "2021-01-12 21:55:00,997: validation_start_date: 1980-10-01 00:00:00\n",
      "2021-01-12 21:55:00,998: finetune_modules: ['head', 'lstm']\n",
      "2021-01-12 21:55:00,999: base_run_dir: /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010\n",
      "2021-01-12 21:55:01,001: is_finetuning: True\n",
      "2021-01-12 21:55:01,002: ### Device cuda:0 will be used for training\n",
      "2021-01-12 21:55:01,030: Starting training from checkpoint /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_1201_213010/model_epoch003.pt\n",
      "2021-01-12 21:55:01,086: Loading basin data into xarray data set.\n",
      "100%|██████████| 1/1 [00:00<00:00, 10.51it/s]\n",
      "2021-01-12 21:55:01,187: Calculating target variable stds per basin\n",
      "100%|██████████| 1/1 [00:00<00:00, 204.34it/s]\n",
      "2021-01-12 21:55:01,200: Create lookup table and convert to pytorch tensor\n",
      "100%|██████████| 1/1 [00:00<00:00, 33.07it/s]\n",
      "# Epoch 1: 100%|██████████| 13/13 [00:01<00:00, 12.81it/s, Loss: 0.0090]\n",
      "2021-01-12 21:55:02,478: Epoch 1 average loss: 0.01336982180006229\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  2.51it/s]\n",
      "2021-01-12 21:55:02,889: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch001/validation_results.p\n",
      "2021-01-12 21:55:02,891:  -- Median validation metrics:NSE: 0.72820\n",
      "2021-01-12 21:55:02,893: Setting learning rate to 5e-05\n",
      "# Epoch 2: 100%|██████████| 13/13 [00:00<00:00, 13.35it/s, Loss: 0.0089]\n",
      "2021-01-12 21:55:03,871: Epoch 2 average loss: 0.010349262290849136\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  4.09it/s]\n",
      "2021-01-12 21:55:04,125: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch002/validation_results.p\n",
      "2021-01-12 21:55:04,126:  -- Median validation metrics:NSE: 0.72768\n",
      "# Epoch 3: 100%|██████████| 13/13 [00:00<00:00, 13.15it/s, Loss: 0.0109]\n",
      "2021-01-12 21:55:05,120: Epoch 3 average loss: 0.011365636968268799\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  3.41it/s]\n",
      "2021-01-12 21:55:05,422: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch003/validation_results.p\n",
      "2021-01-12 21:55:05,424:  -- Median validation metrics:NSE: 0.72449\n",
      "# Epoch 4: 100%|██████████| 13/13 [00:00<00:00, 13.14it/s, Loss: 0.0129]\n",
      "2021-01-12 21:55:06,418: Epoch 4 average loss: 0.011065094015346123\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  4.04it/s]\n",
      "2021-01-12 21:55:06,674: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch004/validation_results.p\n",
      "2021-01-12 21:55:06,677:  -- Median validation metrics:NSE: 0.72565\n",
      "# Epoch 5: 100%|██████████| 13/13 [00:01<00:00, 12.63it/s, Loss: 0.0088]\n",
      "2021-01-12 21:55:07,711: Epoch 5 average loss: 0.01116196932987525\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  3.88it/s]\n",
      "2021-01-12 21:55:07,978: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch005/validation_results.p\n",
      "2021-01-12 21:55:07,979:  -- Median validation metrics:NSE: 0.73648\n",
      "# Epoch 6: 100%|██████████| 13/13 [00:00<00:00, 13.60it/s, Loss: 0.0095]\n",
      "2021-01-12 21:55:08,941: Epoch 6 average loss: 0.010371442096164593\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  4.02it/s]\n",
      "2021-01-12 21:55:09,202: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch006/validation_results.p\n",
      "2021-01-12 21:55:09,204:  -- Median validation metrics:NSE: 0.72587\n",
      "# Epoch 7: 100%|██████████| 13/13 [00:00<00:00, 13.21it/s, Loss: 0.0123]\n",
      "2021-01-12 21:55:10,192: Epoch 7 average loss: 0.010512643350431552\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  4.02it/s]\n",
      "2021-01-12 21:55:10,454: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch007/validation_results.p\n",
      "2021-01-12 21:55:10,456:  -- Median validation metrics:NSE: 0.73680\n",
      "# Epoch 8: 100%|██████████| 13/13 [00:00<00:00, 13.07it/s, Loss: 0.0091]\n",
      "2021-01-12 21:55:11,454: Epoch 8 average loss: 0.009106523858813139\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  3.72it/s]\n",
      "2021-01-12 21:55:11,731: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch008/validation_results.p\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-12 21:55:11,733:  -- Median validation metrics:NSE: 0.74184\n",
      "# Epoch 9: 100%|██████████| 13/13 [00:00<00:00, 13.33it/s, Loss: 0.0075]\n",
      "2021-01-12 21:55:12,714: Epoch 9 average loss: 0.009976885388963498\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  3.84it/s]\n",
      "2021-01-12 21:55:12,985: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch009/validation_results.p\n",
      "2021-01-12 21:55:12,987:  -- Median validation metrics:NSE: 0.72402\n",
      "# Epoch 10: 100%|██████████| 13/13 [00:00<00:00, 13.33it/s, Loss: 0.0094]\n",
      "2021-01-12 21:55:13,966: Epoch 10 average loss: 0.00913747357061276\n",
      "# Validation: 100%|██████████| 1/1 [00:00<00:00,  4.10it/s]\n",
      "2021-01-12 21:55:14,219: Stored results at /home/frederik/Projects/neuralhydrology/examples/06-Finetuning/runs/cudalstm_maurer_531_basins_finetuned_1201_215500/validation/model_epoch010/validation_results.p\n",
      "2021-01-12 21:55:14,222:  -- Median validation metrics:NSE: 0.72906\n"
     ]
    }
   ],
   "source": [
    "finetune(Path(\"finetune.yml\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looking at the validation result, we can see an increase of roughly 0.1 NSE.\n",
    "\n",
    "Last but not least, we will compare the pre-trained and the finetuned model on the test period. For this, we will make use of the `eval_run` function from `neuralhydrolgy.nh_run`. Alternatively, you could evaluate both runs from the command line by calling\n",
    "\n",
    "```bash\n",
    "nh-run evaluate --run-dir /path/to/run_directory/\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-12 21:55:30,573: Using the model weights from runs/cudalstm_maurer_531_basins_1201_213010/model_epoch003.pt\n",
      "# Evaluation: 100%|██████████| 531/531 [03:31<00:00,  2.51it/s]\n",
      "2021-01-12 21:59:02,224: Stored results at runs/cudalstm_maurer_531_basins_1201_213010/test/model_epoch003/test_results.p\n"
     ]
    }
   ],
   "source": [
    "eval_run(run_dir, period=\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we check for the full name of the finetuning run (which we could also extract from the log output above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cudalstm_maurer_531_basins_1201_213010\r\n",
      "cudalstm_maurer_531_basins_finetuned_1201_215500\r\n"
     ]
    }
   ],
   "source": [
    "!ls runs/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can call the `eval_run()` function as above, but pointing to the directory of the finetuned run. By default, this function evaluates the last checkpoint, which can be changed with the `epoch` argument. Here however, we use the default."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-01-12 22:04:31,940: Using the model weights from runs/cudalstm_maurer_531_basins_finetuned_1201_215500/model_epoch010.pt\n",
      "# Evaluation: 100%|██████████| 1/1 [00:00<00:00,  2.43it/s]\n",
      "2021-01-12 22:04:32,361: Stored results at runs/cudalstm_maurer_531_basins_finetuned_1201_215500/test/model_epoch010/test_results.p\n"
     ]
    }
   ],
   "source": [
    "finetune_dir = Path(\"runs/cudalstm_maurer_531_basins_finetuned_1201_215500\")\n",
    "eval_run(finetune_dir, period=\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at the test period results of the pre-trained base model and the finetuned model for the basin that we chose above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Basin 04122200 base model performance: 0.602\n",
      "Performance after finetuning: 0.717\n"
     ]
    }
   ],
   "source": [
    "# load test results of the base run\n",
    "with open(run_dir / \"test/model_epoch003/test_results.p\", \"rb\") as fp:\n",
    "    base_model_results = pickle.load(fp)\n",
    "    \n",
    "# load test results of the finetuned model\n",
    "with open(finetune_dir / \"test/model_epoch010/test_results.p\", \"rb\") as fp:\n",
    "    finetuned_results = pickle.load(fp)\n",
    "    \n",
    "# extract basin performance\n",
    "base_model_nse = base_model_results[basin]['1D']['NSE']\n",
    "finetune_nse = finetuned_results[basin][\"1D\"][\"NSE\"]\n",
    "print(f\"Basin {basin} base model performance: {base_model_nse:.3f}\")\n",
    "print(f\"Performance after finetuning: {finetune_nse:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we see roughly the same performance increase in the test period, which is great. However, note that a) our base model was not optimally trained (we stopped quite early) but also b) the finetuning settings were chosen rather randomly. From our experience so far, you can almost always get performance increases for individual basins with finetuning, but it is difficult to find settings that are universally applicable. However, this tutorial was just a showcase of how easy it actually is to finetune models with the `neuralHydrology` library. Now it is up to you to experiment with it."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
